package preprocess

import (
	"log"
	"strings"

	"github.com/azhao1981/opencc"
	"github.com/go-ego/gse"
	"github.com/longbridge/autocorrect"
	"github.com/mozillazg/go-pinyin"
)

type Config struct {
	Cut bool
}

type Preprocessor struct {
	Config         *Config
	seg            *gse.Segmenter
	s2t            *opencc.OpenCC
	t2s            *opencc.OpenCC
	autocorrectOps []autocorrect.Option
	pinYinOps      *pinyin.Args
}

func NewPreprocessor(config *Config) *Preprocessor {
	pre := &Preprocessor{
		Config: config,
	}
	var seg gse.Segmenter
	seg.LoadDict()
	pre.seg = &seg
	s2t, err := opencc.New("s2t")
	if err != nil {
		log.Fatal(err)
	}
	pre.s2t = s2t
	t2s, err := opencc.New("t2s")
	if err != nil {
		log.Fatal(err)
	}
	pre.t2s = t2s
	pre.autocorrectOps = []autocorrect.Option{}
	pre.pinYinOps = &pinyin.Args{}
	pre.pinYinOps.Style = pinyin.Tone
	pre.pinYinOps.Fallback = func(r rune, a pinyin.Args) []string {
		return []string{string(r)}
	}
	return pre
}

var PreprocessorDefault = NewPreprocessor(&Config{
	Cut: true,
})

func Cut(text string) string {
	words := PreprocessorDefault.seg.Cut(text)
	ret := strings.Builder{}
	for _, word := range words {
		ret.WriteString(word)
		ret.WriteString(" ")
	}
	return strings.TrimSpace(ret.String())
}

func Autocorrect(text string) string {
	return autocorrect.Format(text, PreprocessorDefault.autocorrectOps...)
}

func OpenccT2S(text string) string {
	ret, err := PreprocessorDefault.t2s.Convert(text)
	if err != nil {
		log.Fatal(err)
	}
	return ret
}

func OpenccS2T(text string) string {
	ret, err := PreprocessorDefault.s2t.Convert(text)
	if err != nil {
		log.Fatal(err)
	}
	return ret
}

func Pinyin(text string) string {
	words := pinyin.Pinyin(text, *PreprocessorDefault.pinYinOps)
	ret := strings.Builder{}
	for _, word := range words {
		for _, w := range word {
			ret.WriteString(w)
			ret.WriteString(" ")
		}
	}
	return strings.TrimSpace(ret.String())
}
